package com.example.rsa;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.util.IdUtil;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.core.oidc.endpoint.OidcParameterNames;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;

import java.util.*;
import java.util.stream.Collectors;

public class RsaKeyPairJWKTokenCustomizer implements  OAuth2TokenCustomizer<JwtEncodingContext> {


    @Override
    public void customize(JwtEncodingContext context) {
        Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
        Collection<? extends GrantedAuthority> authorities = authentication.getAuthorities();
        JwtClaimsSet.Builder claims = context.getClaims();
        Object principal = context.getPrincipal().getPrincipal();
        if (principal instanceof UserDetails user) {
            List<String> collect = user.getAuthorities().stream().map(GrantedAuthority::getAuthority).distinct().collect(Collectors.toList());
            claims.claim("roles", collect);
//            for (String s : collect) {
//                attributes.put("authority", s);
//            }

            Set<String> authorizedScopes = context.getAuthorizedScopes();
            collect.addAll(authorizedScopes);
            claims.claim("authorities", collect);
        }
        //自定义一些东西
        claims.claim("id", IdUtil.fastSimpleUUID());
        claims.claim("time", DateUtil.format(new Date(), "yyyy-MM-dd HH:mm:ss"));
        claims.claim("author", "596767880@qq.com");
        claims.claim("application", "spring-security-oauth2");

        if (context.getTokenType().equals(OAuth2TokenType.ACCESS_TOKEN)) {
            // Customize headers/claims for access_token

        } else if (context.getTokenType().getValue().equals(OidcParameterNames.ID_TOKEN)) {
            // Customize headers/claims for id_token

        }
    }
}